from ultralytics import YOLO
from tasks import *
from utils import *
from rich import print


def val(args):
    remove_cache()

    model = YOLO(args.model)
    task = args.task
    if task in ("gen1", "gen4"):
        model.TrainerClass = GenXTrainer
        model.ValidatorClass = GenXValidator
    elif task in ("dsec",):
        model.TrainerClass = DSECTrainer
        model.ValidatorClass = DSECValidator
    elif task in ("sod"):
        model.TrainerClass = SODTrainer
        model.ValidatorClass = SODValidator

    print(model.model)
    model.val(data = {"name": task, "path": args.data, "seq_len": args.t}, batch = args.b)


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("-task", type = str, required = True, help = "The task for training and evaluation.")
    parser.add_argument("-data", type = str, required = True, help = "Path of the data for training and validation.")
    parser.add_argument("-model", type = str, default = "./config/yolov8n.yaml", help = "Path of the model or model config.")
    parser.add_argument("-b", type = int, default = 16, help = "Batch size.")
    parser.add_argument("-t", type = int, default = 5, help = "Time steps.")
    args = parser.parse_args()
    print(args)

    val(args)